from data_resource.data_bases import IndexDaily, engine, MacroBondCN, MacroBondUS
from utilities.utilities_func import get_session
from sqlalchemy import desc
import pandas as pd


def get_index(index_code: list, period: int = 252):
    """获取区间内指数数据，并计算累计收益率"""
    period = len(index_code) * period
    with get_session(engine) as session:
        index_data = session.query(
            IndexDaily.ts_code, IndexDaily.trade_date, IndexDaily.close, IndexDaily.vol
        ).filter(
            IndexDaily.ts_code.in_(index_code)).order_by(desc(IndexDaily.trade_date)).limit(period).all()
        index_data = pd.DataFrame(index_data)

    index_data.sort_values('trade_date', ascending=True, inplace=True)

    index_data['cum_return'] = index_data.groupby('ts_code')['close'].transform(
        lambda x: (x / x.shift(1)).cumprod() - 1)

    return index_data


def get_bondyield(period: int = 10, start_date: str = '2010-01-01', end_date: str = '2023-01-01'):
    """
        中美国债收益率差
    """
    assert period in [1, 10]

    if period == 10:
        with get_session(engine) as session:
            bond_cn = session.query(MacroBondCN.trade_date, MacroBondCN.yield_rate).filter(
                MacroBondCN.trade_date >= start_date,
                MacroBondCN.trade_date <= end_date,
                MacroBondCN.curve_term == 10,
                MacroBondCN.ts_code == '1001.CB',
                MacroBondCN.curve_type == '1'
            ).order_by(MacroBondCN.trade_date).all()
            bond_cn = pd.DataFrame(bond_cn)
            bond_us = session.query(MacroBondUS.date, MacroBondUS.y10).filter(
                MacroBondUS.date >= start_date,
                MacroBondUS.date <= end_date
            ).order_by(MacroBondUS.date).all()
            bond_us = pd.DataFrame(bond_us)

        _data = bond_cn.merge(bond_us, left_on='trade_date', right_on='date', how='inner')
        _data = _data[['trade_date', 'yield_rate', 'y10']]
        _data.rename(columns={'yield_rate': '中债10年', 'y10': '美债10年'}, inplace=True)
        _data['利差'] = _data['中债10年'] - _data['美债10年']

    else:
        with get_session(engine) as session:
            bond_cn = session.query(MacroBondCN.trade_date, MacroBondCN.yield_rate).filter(
                MacroBondCN.trade_date >= start_date,
                MacroBondCN.trade_date <= end_date,
                MacroBondCN.curve_term == 1,
                MacroBondCN.ts_code == '1001.CB',
                MacroBondCN.curve_type == '1'
            ).order_by(MacroBondCN.trade_date).all()
            bond_cn = pd.DataFrame(bond_cn)
            bond_us = session.query(MacroBondUS.date, MacroBondUS.y1).filter(
                MacroBondUS.date >= start_date,
                MacroBondUS.date <= end_date
            ).order_by(MacroBondUS.date).all()
            bond_us = pd.DataFrame(bond_us)

        _data = bond_cn.merge(bond_us, left_on='trade_date', right_on='date', how='inner')
        _data = _data[['trade_date', 'yield_rate', 'y1']]
        _data.rename(columns={'yield_rate': '中债1年', 'y1': '美债1年'}, inplace=True)
        _data['利差'] = _data['中债1年'] - _data['美债1年']

    return _data


if __name__ == '__main__':
    data1 = get_bondyield(period=10, start_date='2010-01-01', end_date='2025-10-13')
    print(data1)
